# ai处理逻辑
import requests
import re
from bs4 import BeautifulSoup
from typing import Optional, List
from .utils import load_config


class HTMLProcessor:
    @staticmethod
    def extract_important_content(html_content: str) -> str:
        """
        提取 HTML 内容中的关键部分。
        :param html_content: 原始 HTML 内容
        :return: 提取后的文本内容
        """
        soup = BeautifulSoup(html_content, "html.parser")
        # 移除脚本和样式标签
        for script in soup(["script", "style"]):
            script.extract()
        # 提取纯文本内容
        return soup.get_text(separator=" ")

    @staticmethod
    def compress_text(text_content: str) -> str:
        """
        压缩提取后的文本内容，移除不必要的空格和换行符。
        :param text_content: 提取后的文本内容
        :return: 压缩后的文本内容
        """
        # 移除多余的空格和换行符
        return re.sub(r'\s+', ' ', text_content).strip()

    @staticmethod
    def chunk_text(text: str, chunk_size: int = 32768) -> List[str]:
        """
        将文本分块。
        :param text: 需要分块的文本
        :param chunk_size: 每块的大小（默认 32k 字符）
        :return: 分块后的文本列表
        """
        return [text[i:i + chunk_size] for i in range(0, len(text), chunk_size)]


class AIClient:
    def __init__(self, api_key: str, api_url: str):
        """
        初始化 AI 客户端。
        :param api_key: API 密钥
        :param api_url: API URL
        """
        self.api_key = api_key
        self.api_url = api_url

    def process_text_chunk(self, chunk: str, fields: str, model: str) -> Optional[str]:
        """
        使用 AI 模型处理文本块。
        :param chunk: 需要处理的文本块
        :param fields: 用户自定义字段（如 "标题, 链接, 补充内容"）
        :param model: 模型名称
        :return: AI 处理后的结果（如果失败则返回 None）
        """
        # 将 fields 参数格式化为提示信息中的一部分
        fields_description = ", ".join(fields.split(","))  # 将逗号分隔的字段转换为自然语言描述
        prompt = (
            f"对以下文本内容进行总结和关键信息提取，将与搜索关键词有关内容整理为表格，包含：{fields_description}。"
            f"若没有相关信息则直接跳过，不输出任何内容: {chunk}"
        )

        payload = {
            "model": model,
            "messages": [
                {
                    "role": "user",
                    "content": prompt
                }
            ],
            "max_tokens": 1000
        }
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }

        try:
            response = requests.post(self.api_url, headers=headers, json=payload)
            response.raise_for_status()
            return response.json()['choices'][0]['message']['content']
        except requests.exceptions.RequestException as e:
            print(f"AI 处理失败: {e}")
            return None


def process_with_ai(html_content: str, model: str, fields: Optional[str] = None) -> Optional[str]:
    """
    使用 AI 模型处理页面 HTML 内容。
    :param html_content: 页面的 HTML 内容
    :param model: 模型名称
    :param fields: 用户自定义字段（如 "标题, 链接, 补充内容"），默认为 None
    :return: AI 处理后的数据（如果失败则返回 None）
    """
    config = load_config('../config/config.yml')
    ai_config = config.get('siliconflow', {})
    if not all(key in ai_config for key in ['api_key', 'api_url']):
        print("配置文件缺少必要的 AI 配置项")
        return None

    ai_client = AIClient(
        api_key=ai_config['api_key'],
        api_url=ai_config['api_url']
    )

    # 提取并压缩 HTML 内容
    text_content = HTMLProcessor.extract_important_content(html_content)
    compressed_content = HTMLProcessor.compress_text(text_content)
    chunks = HTMLProcessor.chunk_text(compressed_content)

    # 处理每个文本块
    processed_results = []
    for chunk in chunks:
        result = ai_client.process_text_chunk(chunk, model, fields)  # 传递 fields 参数
        if result:
            processed_results.append(result)

    return " ".join(processed_results) if processed_results else None
